package com.zzh.partnersys.ai.vectorstore;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjUtil;
import com.zzh.partnersys.ai.entity.dto.AIDocumentDTO;
import com.zzh.partnersys.ai.util.DimensionReducerUtil;
import com.zzh.partnersys.ai.util.ESRestClientUtil;
import com.zzh.partnersys.ai.util.ThreadLocalContext;
import com.zzh.partnersys.common.constant.ElasticSearchConstant;
import jakarta.annotation.PostConstruct;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.*;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author: zzh
 * @date: 2025/11/16 15:11:10
 * @version: 1.0
 */
@Component
@AllArgsConstructor
@Slf4j
public class ElasticSearchVectorStore implements VectorStore {

    private final ESRestClientUtil esRestClientUtil;

    private final EmbeddingModel dashscopeEmbeddingModel;

    /**
     * 初始化构建，可以从加载数据到ES中存储
     */
    @PostConstruct
    private void init(){

    }

    @Override
    public void add(List<Document> documents) {
        if (CollUtil.isEmpty(documents)){
            return;
        }
        //批量生成Embedding（DashScope单次最多支持25条文本）
        final int batchSize = 25;
        List<AIDocumentDTO> aiDocumentDTOS = new ArrayList<>(documents.size());
        for (int batchStart = 0; batchStart < documents.size(); batchStart += batchSize) {
            int batchEnd = Math.min(batchStart + batchSize, documents.size());
            List<Document> batchDocuments = documents.subList(batchStart, batchEnd);
            List<String> batchTexts = batchDocuments.stream().map(Document::getText).toList();

            List<float[]> vectorsList = dashscopeEmbeddingModel.embed(batchTexts);

            for (int j = 0; j < batchDocuments.size(); j++) {
                Document document = batchDocuments.get(j);
                Map<String, Object> metadata = document.getMetadata();
                float[] floats = vectorsList.get(j);
                AIDocumentDTO aiDocumentDTO = new AIDocumentDTO();
                aiDocumentDTO.setId(document.getId());
                aiDocumentDTO.setContent(document.getText());
                // 验证向量维度（应该是1536）
                aiDocumentDTO.setVectors(DimensionReducerUtil.reduceDims(floats, 1024));
            
            // 从元数据中获取文档信息
            // 优先使用 docUrl/docName，如果没有则使用 sourceUrl/sourceName
            String docUrl = (String) metadata.get("docUrl");
            String docName = (String) metadata.get("docName");
            
            // 处理创建时间，如果为null则使用当前时间
            aiDocumentDTO.setCreateTime((String)metadata.get("createTime"));

            // 设置新增的元数据字段
            Object userIdObj = metadata.get("userId");
            if (userIdObj != null) {
                if (userIdObj instanceof Long) {
                    aiDocumentDTO.setUserId((Long) userIdObj);
                } else if (userIdObj instanceof Number) {
                    aiDocumentDTO.setUserId(((Number) userIdObj).longValue());
                }
            }
            Object assistantIdObj = metadata.get("assistantId");
            if (assistantIdObj != null) {
                if (assistantIdObj instanceof Long) {
                    aiDocumentDTO.setAssistantId((Long) assistantIdObj);
                } else if (assistantIdObj instanceof Number) {
                    aiDocumentDTO.setAssistantId(((Number) assistantIdObj).longValue());
                }
            }
            aiDocumentDTO.setDocName(docName);
            aiDocumentDTO.setDocUrl(docUrl);
            aiDocumentDTO.setDocType((String) metadata.get("docType"));
            
            Object chunkIndexObj = metadata.get("chunkIndex");
            if (chunkIndexObj != null) {
                if (chunkIndexObj instanceof Integer) {
                    aiDocumentDTO.setChunkIndex((Integer) chunkIndexObj);
                } else if (chunkIndexObj instanceof Number) {
                    aiDocumentDTO.setChunkIndex(((Number) chunkIndexObj).intValue());
                }
            }
            
            Object totalChunksObj = metadata.get("totalChunks");
            if (totalChunksObj != null) {
                if (totalChunksObj instanceof Integer) {
                    aiDocumentDTO.setTotalChunks((Integer) totalChunksObj);
                } else if (totalChunksObj instanceof Number) {
                    aiDocumentDTO.setTotalChunks(((Number) totalChunksObj).intValue());
                }
            }
            
                aiDocumentDTOS.add(aiDocumentDTO);
            }
        }
        
        // 从第一个文档的元数据中获取 assistantId，如果没有则尝试从 ThreadLocalContext 获取
        Long assistantId = null;
        Long userId = null;
        
        if (!documents.isEmpty()) {
            Map<String, Object> firstMetadata = documents.get(0).getMetadata();
            Object assistantIdObj = firstMetadata.get("assistantId");
            if (assistantIdObj != null) {
                if (assistantIdObj instanceof Long) {
                    assistantId = (Long) assistantIdObj;
                } else if (assistantIdObj instanceof Number) {
                    assistantId = ((Number) assistantIdObj).longValue();
                }
            }
        }
        
        // 如果 metadata 中没有，尝试从 ThreadLocalContext 获取
        if (assistantId == null) {
            try {
                assistantId = ThreadLocalContext.getAssistantId();
            } catch (IllegalStateException e) {
                // ThreadLocalContext 未初始化，使用默认值
                log.warn("无法从 ThreadLocalContext 获取 assistantId，使用默认值");
            }
        }
        
        try {
            userId = ThreadLocalContext.getUserId();
        } catch (IllegalStateException e) {
            // ThreadLocalContext 未初始化，使用默认值
            log.warn("无法从 ThreadLocalContext 获取 userId，使用默认值");
        }
        
        // 如果都获取不到，使用默认值（兼容旧代码）
        if (userId == null) {
            userId = 0L;
        }
        if (assistantId == null) {
            assistantId = 0L;
        }
        
        //转换Document 为  AIDocumentDTO
        esRestClientUtil.insertVectorBatchDocument(aiDocumentDTOS, ElasticSearchConstant.DOCUMENT_INDEX_NAME);
    }

    @Override
    public void delete(List<String> idList) {
        if (CollUtil.isEmpty(idList)){
            return;
        }
        esRestClientUtil.deleteVectorBatchDocument(idList, String.format(ElasticSearchConstant.DOCUMENT_INDEX_NAME, "user_id", "assistant_id"));
    }

    @Override
    public void delete(Filter.Expression filterExpression) {
        //todo 实现删除根据表达式删除
    }

    @NotNull
    @Override
    public List<Document> similaritySearch(@NotNull SearchRequest request) {
        if (ObjUtil.isEmpty(request)){
            return List.of();
        }
        List<Double> queryVector = convertToDoubleListEmbedding(request.getQuery());
        //向量转换
        int topK = request.getTopK();
        double similarityThreshold = request.getSimilarityThreshold(); // 保留用于后续相似度过滤
        Long userId = ThreadLocalContext.getUserId();
        Long assistantId = ThreadLocalContext.getAssistantId();
        ThreadLocalContext.clearContext();
        System.out.println("AI回答用户ID：" + userId);
        System.out.println("AI回答助手ID：" + assistantId);

        List<AIDocumentDTO> aiDocumentDTOS = esRestClientUtil.knnVectorSearch(queryVector,
                ElasticSearchConstant.DOCUMENT_INDEX_NAME, topK, userId, assistantId,similarityThreshold);
        //转换为文档
        List<Document> documents = new ArrayList<>();
        for (AIDocumentDTO aiDocumentDTO : aiDocumentDTOS) {
            HashMap<String, Object> metadata = new HashMap<>();
            metadata.put("createTime", aiDocumentDTO.getCreateTime());
            // 添加新增的元数据字段
            if (aiDocumentDTO.getAssistantId() != null) {
                metadata.put("assistantId", aiDocumentDTO.getAssistantId());
            }
            if (aiDocumentDTO.getUserId() != null){
                metadata.put("userId", aiDocumentDTO.getUserId());
            }
            if (aiDocumentDTO.getDocName() != null) {
                metadata.put("docName", aiDocumentDTO.getDocName());
            }
            if (aiDocumentDTO.getDocUrl() != null) {
                metadata.put("docUrl", aiDocumentDTO.getDocUrl());
            }
            if (aiDocumentDTO.getDocType() != null) {
                metadata.put("docType", aiDocumentDTO.getDocType());
            }
            if (aiDocumentDTO.getChunkIndex() != null) {
                metadata.put("chunkIndex", aiDocumentDTO.getChunkIndex());
            }
            if (aiDocumentDTO.getTotalChunks() != null) {
                metadata.put("totalChunks", aiDocumentDTO.getTotalChunks());
            }

            Document document = new Document(aiDocumentDTO.getId(), aiDocumentDTO.getContent(), metadata);
            documents.add(document);
        }
        return documents;
    }


    private List<Double> convertToDoubleListEmbedding(String query) {
        EmbeddingOptions embeddingOptions = EmbeddingOptionsBuilder.builder()
                .withModel("text-embedding-v2")
                .withDimensions(1536)
                .build();
        EmbeddingResponse call = dashscopeEmbeddingModel.call(new EmbeddingRequest(List.of(query), embeddingOptions));
        List<Embedding> results = call.getResults();
        float[] output = results.get(0).getOutput();
        
        // 将 1536 维向量降维到 1024 维（与 ES mapping 中的维度一致）
        float[] reducedVector = DimensionReducerUtil.reduceDims(output, 1024);
        
        // 转换为 List<Double>
        List<Double> queryVector = new ArrayList<>(1024);
        for (float v : reducedVector) {
            queryVector.add((double) v);
        }
        return queryVector;
    }
}
